package cn.njupt.map;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;

public class TreeMap<K,V> implements Map<K,V>{
    public static final boolean RED=true;
    public static final boolean BLACK=false;

    private Comparator<K> comparator;

    private int size;
    private Node<K,V> root;

    public TreeMap(Comparator<K> comparator){
        this.comparator=comparator;
    }
    public TreeMap(){
        this(null);
    }
    @Override
    public int size() {
        return size;
    }

    @Override
    public boolean isEmpty() {
        return size==0;
    }

    @Override
    public void clear() {
        size=0;
        root=null;
    }

    public void checkKeyIsNull(K key){
        if(key==null){
            throw new IllegalArgumentException("传递的参数为空！");
        }
    }

    @Override
    public V put(K key, V value) {
        checkKeyIsNull(key);
        if(root==null){
            Node<K,V> node=new Node(key,value,null);
            root=node;
            size++;
            afterPut(node);
            return null;
        }else{
            Node<K,V> cur=root;
            Node<K,V> parent=null;
            int cmp=0;
            while(cur!=null){
                cmp=compare(cur.key,key);
                parent=cur;
                if(cmp>0){
                    //说明在左子树
                    cur=cur.left;
                }else if(cmp<0){
                    //说明在右子树
                    cur=cur.right;
                }else{
                    //相等
                    V oldValue=cur.value;
                    cur.value=value;
                    cur.key=key;
                    return oldValue;
                }
            }
            Node<K,V> newNode=new Node<>(key,value,parent);
            if(cmp>0){
                //左子树
                parent.left=newNode;
            }else{
                //右子树
                parent.right=newNode;
            }
            afterPut(newNode);
        }
        size++;
        return null;
    }


    @Override
    public V get(K key) {
        return node(key)!=null?node(key).value:null;
    }

    @Override
    public V remove(K key) {
        return remove(node(key));
    }

    public V remove(Node<K,V> node) {
        if(node==null)return null;

        V oldValue=node.value;

        if(node.hasTwoChildren()){
            //拿到后继节点进行替换
            Node<K,V> replacement=getSuccessor(node);
            node.key=replacement.key;
            node.value=replacement.value;
            //将node指向replacement,后期只要删除node即可
            node=replacement;
        }
        Node<K,V> replacement=node.left==null?node.right:node.left;
        if(replacement!=null){
            //度为1
            if(node.parent==null){
                //根节点
                root=replacement;
                root.parent=null;
            }else if(node.parent.left==node){
                //被删除节点相对于根节点是左子树
                node.parent.left=replacement;
                replacement.parent=node.parent;
            }else{
                //被删除节点相对于根节点是右子树
                node.parent.right=replacement;
                replacement.parent=node.parent;
            }
            afterRemove(replacement);
        }else if(node.parent==null){
            //==null说明是叶子节点
            root = null;
        }else{
            if(node.parent.left==node){
                //左子树
                node.parent.left=null;
            }else{
                node.parent.right=null;
            }
            afterRemove(node);
        }
        size--;
        return oldValue;
    }

    @Override
    public boolean containsKey(K key) {
        return node(key)!=null;
    }

    @Override
    public boolean containsValue(V value) {
        //利用层序遍历来判断是否包含value
        if(root==null)return false;
        Queue<Node<K,V>> queue=new LinkedList<>();
        queue.offer(root);
        while(!queue.isEmpty()){
            Node<K,V> curNode=queue.poll();
            if(isEquals(curNode.value,value)){
                return true;
            }
            if(curNode.left!=null){
                queue.offer(curNode.left);
            }
            if(curNode.right!=null){
                queue.offer(curNode.right);
            }
        }
        return false;
    }

    public boolean isEquals(V v1,V v2){
        return v1==null?v2==null:v1.equals(v2);
    }

    @Override
    public void traversal(Visitor<K, V> visitor) {
        if(visitor==null)return;
        traversal(root,visitor);
    }

    public void traversal(Node<K,V> node,Visitor<K,V> visitor){
        if(node==null || visitor.stop)return;
        traversal(node.left,visitor);
        if(visitor.stop)return;
        visitor.stop=visitor.visitor(node.key, node.value);
        traversal(node.right,visitor);
    }

    public Node<K,V> node(K key){
        checkKeyIsNull(key);
        Node<K,V> cur=root;
        while(cur!=null){
            int cmp=compare(cur.key,key);
            if(cmp>0){
                //说明在左边
                cur=cur.left;
            }else if(cmp<0){
                //说明在右边
                cur=cur.right;
            }else{
                //说明相等
                return cur;
            }
        }
        return null;
    }

    /**
     * 得到node的后继节点
     * @param node 传递的当前节点
     * @return
     */
    public Node<K,V> getSuccessor(Node<K,V> node){
        if(node==null) return node;
        if(node.right!=null){
            return getLeftNode(node.right);
        }else{
            Node<K,V> parent=node.parent;
            while(parent!=null && parent.left!=node){
                node=parent;
                parent=node.parent;

            }
            return parent;
        }
    }
    public Node<K,V> getLeftNode(Node<K,V> node){
        while (node.left!=null){
            node=node.left;
        }
        return node;
    }

    public int compare(K e1,K e2){
        if(comparator!=null){
            return comparator.compare(e1,e2);
        }
        return ((Comparable<K>)e1).compareTo(e2);
    }


    public void afterPut(Node<K,V> child) {
        Node<K,V> parent=child.parent;
        if(parent==null){
            colorBlack(child);
            return;
        }
        if(isBlack(parent))return;
        Node<K,V> grand=parent.parent;
        Node<K,V> uncle=parent.brother();

        if(isRed(uncle)){
            //叔父节点是红色
            colorBlack(parent);
            colorBlack(uncle);
            colorRed(parent.parent);
            afterPut(parent.parent);
            return;
        }
        //叔父节点不是红色，就要分情况进行LL,LR,RR,RL
        if(parent.isLeftChild()){
            //L
            if(child.isLeftChild()){
//                LL
                colorBlack(parent);
                colorRed(grand);
                rotateRight(grand);
            }else{
                //LR
                colorBlack(child);
                colorRed(grand);
                rotateLeft(parent);
                rotateRight(grand);
            }
        }else{
            if(child.isRightChild()){
//                RR
                colorRed(grand);
                colorBlack(parent);
                rotateLeft(grand);
            }else{
//                RL
                colorBlack(child);
                colorRed(grand);
                rotateRight(parent);
                rotateLeft(grand);
            }
        }
    }

    //判断是否是红色节点
    public boolean isRed(Node<K,V> node){
        return node==null?false:node.color==RED;
    }
    //判断是否是黑色节点
    public boolean isBlack(Node<K,V> node){
        return node==null?true:node.color==BLACK;
    }
    public Node<K,V> color(Node<K,V> node,boolean color){
        if(node==null)return node;
        node.color=color;
        return node;
    }

    public boolean colorOf(Node<K,V> node){
        return node==null?BLACK:node.color;
    }

    //染成红色
    public void colorRed(Node<K,V> node){
        color(node,RED);
    }
    //染成黑色
    public void colorBlack(Node<K,V> node){
        color(node,BLACK);
    }

    /**
     * 左旋转
     * @param grand 传递的最低不平衡的那个节点
     */
    public void rotateLeft(Node<K,V> grand){
        Node<K,V> parent=grand.right;
        Node<K,V> child=parent.left;

        grand.right=child;

        parent.left=grand;

        afterRotate( grand,  parent,  child);
    }

    protected void afterRotate(Node<K,V> grand, Node<K,V> parent, Node<K,V> child) {
        if (grand.parent == null) {
            root = parent;
        } else if (grand.isLeftChild()) {
            grand.parent.left = parent;
        } else {
            grand.parent.right = parent;
        }
        parent.parent = grand.parent;
        grand.parent = parent;
        //child可能为null
        if (child != null) {
            child.parent = grand;
        }
    }

    /**
     * 右旋转
     * @param grand 传递的最低不平衡的那个节点
     */
    public void rotateRight(Node<K,V> grand){
        Node<K,V> parent=grand.left;
        Node<K,V> child=parent.right;

        grand.left=child;

        parent.right=grand;
        afterRotate(grand, parent, child);
    }
    public void afterRemove(Node<K,V> node) {


        //下面删除的黑色节点，并且黑色节点不是叶子节点
        if(isRed(node)){
            colorBlack(node);
            return;
        }

        Node<K,V> parent=node.parent;
        if(parent==null){
            return;
        }

        //下面删除的是黑色叶子节点
        boolean left=parent.left==null || node.isLeftChild();
        if(left){
            //被删除的节点相当于父节点的左边
            Node<K,V> brother=parent.right;
            if(isRed(brother)){
                colorBlack(brother);
                colorRed(parent);
                rotateLeft(parent);
                brother=parent.right;
            }
            if(isBlack(brother.left) &&  isBlack(brother.right)){
                //兄弟节点是一个黑色节点，并且黑色节点没有红色子节点
                boolean parentBlack=isBlack(parent);
                colorRed(brother);
                colorBlack(parent);
                if(parentBlack){
                    afterRemove(parent);
                }

            }else{
                //兄弟节点是一个黑色节点，并且黑色节点有红色子节点
                if(isBlack(brother.right)){
                    rotateRight(brother);
                    brother=parent.right;
                }
                color(brother,colorOf(parent));
                colorBlack(parent);
                colorBlack(brother.right);
                rotateLeft(parent);
            }
        }else{
            //被删除的节点相当于父节点的右边
            Node<K,V> brother=parent.left;
            if(isRed(brother)){
                colorBlack(brother);
                colorRed(parent);
                rotateRight(parent);
                brother=parent.left;
            }
            if(isBlack(brother.left) &&  isBlack(brother.right)){
                //兄弟节点是一个黑色节点，并且黑色节点没有红色子节点
                boolean parentBlack=isBlack(parent);
                colorRed(brother);
                colorBlack(parent);
                if(parentBlack){
                    afterRemove(parent);
                }

            }else{
                //兄弟节点是一个黑色节点，并且黑色节点有红色子节点
                if(isBlack(brother.left)){
                    rotateLeft(brother);
                    brother=parent.left;
                }
                color(brother,colorOf(parent));
                colorBlack(parent);
                colorBlack(brother.left);
                rotateRight(parent);
            }
        }
    }


    private static class Node<K,V>{
        K key;
        V value;
        Node<K,V> left;
        Node<K,V> right;
        Node<K,V> parent;
        boolean color=RED;


        public Node(K key,V value,Node<K,V> parent){
            this.key=key;
            this.value=value;
            this.parent=parent;
        }

        public boolean hasTwoChildren(){
            return left!=null && right!=null;
        }

        public boolean isLeftChild(){
            return parent!=null && parent.left==this;
        }
        public boolean isRightChild(){
            return parent!=null && parent.right==this;
        }

        public Node<K,V> brother(){
            if(isLeftChild()){
                return parent.right;
            }
            if(isRightChild()){
                return parent.left;
            }
            return null;
        }

    }

}
